package com.microservice.core.easyexcel;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.builder.ExcelWriterBuilder;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.io.OutputStream;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.function.BiFunction;

@Slf4j
public class EasyExcelExportUtil {

    /**
     * 导出数据到HttpServletResponse
     *
     * @param response       HttpServletResponse对象
     * @param fileName       文件名 (纯文件名不含路径)
     * @param sheetMaxRowNum sheet页码最大行数
     * @param querySize      每次查询数据库的数据量，可自行调整，但要小于等于sheetMaxRowNum
     * @param headClass      表头信息类型
     * @param param          查询参数
     * @param dataFunc       数据查询函数
     * @throws IOException 如果发生IO错误
     */
    public static <T, U> void exportToResponse(HttpServletResponse response, String fileName, long sheetMaxRowNum, long querySize, Class<?> headClass, U param, BiFunction<Page<T>, U, Page<T>> dataFunc) throws IOException {
        paramCheck(sheetMaxRowNum, querySize);
        responseSet(response, fileName);
        try (OutputStream outputStream = response.getOutputStream()) {
            exportData(outputStream, sheetMaxRowNum, querySize, headClass, param, dataFunc);
        }
    }

    /**
     * 导出数据到HttpServletResponse，支持多线程并发查询导出，适合大数据量导出
     *
     * @param response       HttpServletResponse对象
     * @param fileName       文件名
     * @param sheetMaxRowNum sheet页码最大行数
     * @param querySize      每次查询数据库的数据量，可自行调整，但要小于等于sheetMaxRowNum
     * @param headClass      表头信息类型
     * @param dataFunc       数据查询函数
     */
    public static <T, U> void concurrentExportToResponse(HttpServletResponse response, String fileName, long sheetMaxRowNum, long querySize, Class<?> headClass, U param, BiFunction<Page<T>, U, Page<T>> dataFunc) throws IOException {
        paramCheck(sheetMaxRowNum, querySize);
        responseSet(response, fileName);
        ExecutorService executorService = Executors.newFixedThreadPool(10);
        try (OutputStream outputStream = response.getOutputStream()) {
            concurrentExportData(executorService, outputStream, sheetMaxRowNum, querySize, headClass, param, dataFunc);
        } finally {
            executorService.shutdown();
        }
    }

    /**
     * 导出数据到Excel文件
     *
     * @param fileName  文件名 (文件的全路径)
     * @param headClass 表头信息
     * @param dataFunc  数据查询函数
     * @param <T>       数据类型
     */
    public static <T, U> void exportToExcel(String fileName, long sheetMaxRowNum, long querySize, Class<?> headClass, U param, BiFunction<Page<T>, U, Page<T>> dataFunc) {
        paramCheck(sheetMaxRowNum, querySize);
        exportData(fileName, sheetMaxRowNum, querySize, headClass, param, dataFunc);
    }

    public static <T, U> void concurrentExportToExcel(String fileName, long sheetMaxRowNum, long querySize, Class<?> headClass, U param, BiFunction<Page<T>, U, Page<T>> dataFunc) {
        paramCheck(sheetMaxRowNum, querySize);
        ExecutorService executorService = getExecutorService();
        concurrentExportData(executorService, fileName, sheetMaxRowNum, querySize, headClass, param, dataFunc);
        executorService.shutdown();
    }

    /**
     * 辅助方法，用于导出数据
     * 本方法支持多sheet页导出，适合大数量导出，分页查询数据库
     *
     * @param s         导出目标对象
     * @param headClass 头信息类型
     * @param dataFunc  数据查询函数
     * @param <T>       数据类型
     */
    private static <S, T, U> void exportData(S s, long sheetMaxRowNum, long querySize, Class<?> headClass, U param, BiFunction<Page<T>, U, Page<T>> dataFunc) {
        ExcelWriter excelWriter = null;
        try {
            ExcelWriterBuilder writerBuilder = getExcelWriterBuilder(s, headClass);
            excelWriter = writerBuilder.build();
            int sheetNo = 1;
            long current = 1;
            long rowsInCurrentSheet = 0;
            String sheetName = "Sheet" + sheetNo;

            // 循环一直查询数据库直到查询不到数据为止
            while (true) {
                Page<T> page = new Page<>(current, querySize);
                page = dataFunc.apply(page, param);
                List<T> records = page.getRecords();

                // 如果当前页没有数据，则结束循环
                if (records.isEmpty()) {
                    break;
                }
                log.info("第{}次写入sheetName:{}", current, sheetName);
                // 写入数据到Sheet页
                excelWriter.write(records, EasyExcel.writerSheet(sheetNo, sheetName).build());

                // 更新分页参数
                current++;
                rowsInCurrentSheet += records.size();

                // 如果当前Sheet页已写满MAX_ROWS_PER_SHEET条数据，则开始新的Sheet页
                if (rowsInCurrentSheet >= sheetMaxRowNum) {
                    rowsInCurrentSheet = 0;
                    sheetNo++;
                    sheetName = "Sheet" + sheetNo;
                }
            }
        } finally {
            // 关闭ExcelWriter
            if (excelWriter != null) {
                excelWriter.finish();
            }
        }
    }

    private static <S, T, U> void concurrentExportData(ExecutorService executorService, S s, long sheetMaxRowNum, long querySize, Class<?> headClass, U param, BiFunction<Page<T>, U, Page<T>> dataFunc) {
        ExcelWriterBuilder writerBuilder = getExcelWriterBuilder(s, headClass);
        try (ExcelWriter excelWriter = writerBuilder.build();) {
            int sheetNo = 1;
            long current = 1;
            long rowsInCurrentSheet = 0;
            String sheetName = "Sheet" + sheetNo;
            boolean continueQuery = true;
            List<CompletableFuture<Page<T>>> futures = new ArrayList<>();
            while (continueQuery) {
                Page<T> page = new Page<>(current, querySize);
                Callable<Page<T>> task = () -> dataFunc.apply(page, param);
                CompletableFuture<Page<T>> future = CompletableFuture.supplyAsync(() -> {
                    try {
                        return task.call();
                    } catch (Exception e) {
                        log.error("并发查询失败", e);
                        return new Page<>(); // 返回空页以继续处理
                    }
                }, executorService);
                futures.add(future);
                current++;

                if (futures.size() >= 10 || current > 100) { // 假设每次并发查询10个分页，或者达到一定数量后处理
                    List<Page<T>> results = futures.stream().map(CompletableFuture::join).toList();

                    for (Page<T> resultPage : results) {
                        List<T> records = resultPage.getRecords();
                        if (records.isEmpty()) {
                            continueQuery = false;
                            break;
                        }
                        excelWriter.write(records, EasyExcel.writerSheet(sheetNo, sheetName).build());
                        rowsInCurrentSheet += records.size();
                        if (rowsInCurrentSheet >= sheetMaxRowNum) {
                            rowsInCurrentSheet = 0;
                            sheetNo++;
                            sheetName = "Sheet" + sheetNo;
                        }
                    }
                    futures.clear();
                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 获取ExcelWriterBuilder
     */
    private static <T> ExcelWriterBuilder getExcelWriterBuilder(T t, Class<?> head) {
        if (t instanceof String) {
            return EasyExcel.write((String) t, head);
        }
        if (t instanceof OutputStream) {
            return EasyExcel.write((OutputStream) t, head);
        }
        return null;
    }

    /**
     * 设置响应头
     */
    private static void responseSet(HttpServletResponse response, String fileName) throws UnsupportedEncodingException {
        // 设置响应头
        response.setContentType("application/vnd.ms-excel");
        response.setCharacterEncoding("utf-8");
        String encodedFileName = URLEncoder.encode(fileName, StandardCharsets.UTF_8).replaceAll("\\+", "%20");
        response.setHeader("Content-disposition", "attachment;filename*=utf-8''" + encodedFileName + ".xlsx");
    }

    /**
     * 获取线程池
     */
    private static ExecutorService getExecutorService() {
        return Executors.newFixedThreadPool(10);
    }

    /**
     * 参数校验
     *
     * @param sheetMaxRowNum sheet最大行数
     * @param querySize      分页查询每页大小
     */
    private static void paramCheck(long sheetMaxRowNum, long querySize) {
        if (querySize > sheetMaxRowNum) {
            throw new RuntimeException("querySize不能大于sheetMaxRowNum");
        }
    }
}
